package wolf.iu.framework;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;
import java.util.TreeMap;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.tree.ClassNode;

public class ClassesHandler {

    private final TreeMap<String, ClassNode> classes = new TreeMap<String, ClassNode>();

    public ClassesHandler(final File f) throws IOException {
        this(new JarInputStream(new FileInputStream(f)));
    }

    public ClassesHandler(final JarInputStream jis) throws IOException {
        JarEntry entry;
        while ((entry = jis.getNextJarEntry()) != null) {
            final String name = entry.getName();
            if (!name.endsWith(".class")) {
                jis.closeEntry();
                continue;
            }
            final ByteArrayOutputStream baos = new ByteArrayOutputStream();
            final byte[] buffer = new byte[1024];
            int a = 0;
            while ((a = jis.read(buffer)) != -1) {
                baos.write(buffer, 0, a);
            }

            final ClassReader cr = new ClassReader(baos.toByteArray());
            final ClassNode cn = new ClassNode();
            cr.accept(cn, ClassReader.SKIP_DEBUG | ClassReader.SKIP_FRAMES);

            //System.out.println(cn.name + " loaded");
            classes.put(cn.name, cn);

            jis.closeEntry();
        }
        jis.close();
    }

    public Set<ClassContainer> getClasses() {
        final HashSet<ClassContainer> set = new HashSet<ClassContainer>();
        for (final ClassNode node : classes.values()) {
            set.add(new ClassContainer(node));
        }
        return set;
    }
    
    public ClassNode getClass(String name) {
        if (name.endsWith(";")) {
            name = name.substring(1, name.length()-1);
        }
        if (!classes.containsKey(name)) {
            System.err.println("Referencing to invalid class " + name);
            return null;
        }
        return this.classes.get(name);
    }

}
